import os
import re
import openpyxl
from openpyxl import Workbook
from openpyxl.styles import PatternFill
from tqdm import tqdm
import logging
from multiprocessing import Pool, cpu_count, current_process

# 设置日志文件名和Excel文件名
log_filename = 'process_asm.log'
xlsx_filename = 'output.xlsx'

# 删除现有的日志文件和Excel文件（如果存在）
if os.path.exists(log_filename):
    os.remove(log_filename)

if os.path.exists(xlsx_filename):
    os.remove(xlsx_filename)

# 配置日志格式
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')

# 定义汇编指令及其时钟周期
instruction_cycles = {
    'MOV': 1,        # Move data
    'ADD': 1,        # Addition
    'SUB': 1,        # Subtraction
    'MUL': 3,        # Multiplication
    'DIV': 12,       # Division
    'NOP': 1,        # No operation
    'JMP': 2,        # Jump
    'CALL': 3,       # Call subroutine
    'RET': 2,        # Return from subroutine
    'PUSH': 2,       # Push onto stack
    'POP': 2,        # Pop from stack
    'CMP': 1,        # Compare
    'AND': 1,        # Logical AND
    'OR': 1,         # Logical OR
    'XOR': 1,        # Logical XOR
    'SHL': 1,        # Shift left
    'SHR': 1,        # Shift right
    'INC': 1,        # Increment
    'DEC': 1,        # Decrement
    'MOVH': 1,       # Move high
    'LEA': 1,        # Load effective address
    'LD': 2,         # Load
    'ST': 2,         # Store
    'JEQ': 2,        # Jump if equal
    'JLEZ': 2,       # Jump if less or equal to zero
    'JNE': 2,        # Jump if not equal
    'JNZ': 2,        # Jump if not zero
    'EXTR': 2,       # Extract
    'SAT': 2,        # Saturate
    'MSUB': 3,       # Multiply and subtract
    'LT': 2,         # Less than
    'JGEZ': 2,       # Jump if greater or equal to zero
    'SHA': 2,        # Shift arithmetic
    'UTOF': 2,       # Unsigned to float
    'FTOIZ': 2,      # Float to integer with zero
    'SH': 2,         # Shift
    'ITOF': 2,       # Integer to float
    'GE': 2,         # Greater or equal
    'JLTZ': 2,       # Jump if less than zero
    'JZ': 2,         # Jump if zero
    'JGE': 2,        # Jump if greater or equal
    'MAX': 2,        # Maximum
    'EQ': 2,         # Equal
    'JLT': 2,        # Jump if less than
    'J': 2,          # Jump
    'NE': 2,         # Not equal
    'ADDI': 2,       # Add immediate
    'LOOP': 2,       # Loop
    'ANDN': 2,       # AND with NOT
    'RSUB': 2,       # Reverse subtract
    'INSERT': 2,     # Insert
    'ADDSC': 2,      # Add with scale
    'JI': 2,         # Jump indirect
    'ORN': 2,        # OR with NOT
    'MADD': 2,       # Multiply and add
    'NOT': 2,        # Logical NOT
    'NOR': 2,        # Logical NOR
    'ADDIH': 2,      # Add immediate high
    'CADD': 2,       # Complex add
    'FTOUZ': 2,      # Float to unsigned zero
    'SEL': 2,        # Select
    'MIN': 2,        # Minimum
    'ABS': 2,        # Absolute
    'CMOV': 2,       # Conditional move
    'ABSDIF': 2,     # Absolute difference
    'SELN': 2,       # Select with negation
    'NAND': 2,       # Logical NAND
    'SHAS': 2,       # Shift arithmetic with saturation
    'CMOVN': 2,      # Conditional move with negation
    'CALLI': 2,      # Call indirect
    'JGTZ': 2,       # Jump if greater than zero
    'DEXTR': 2,      # Double extract
    'DSYNC': 2,      # Data synchronization
    'ISYNC': 2,      # Instruction synchronization
    'MFCR': 2,       # Move from control register
    'RFE': 2,        # Return from exception
    'MTCR': 2,       # Move to control register
    'RSLCX': 2,      # Rotate and shift with carry
    'NEZ': 2,        # Not equal to zero
    'DISABLE': 2,    # Disable
    'SVLCX': 2,      # Save local context
    'CSUB': 2,       # Complex subtract
    'JNED': 2,       # Jump if not equal double
    'LDMST': 2,      # Load multiple and store
    'IMASK': 2,      # Interrupt mask
    'ENABLE': 2,     # Enable
    'LHA': 2,        # Load halfword and add
    'INS': 2,        # Insert
    'SYSCALL': 2,    # System call
    'XNOR': 2,       # Logical XNOR
    'EQZ': 1,        # Equal to zero
    'CADDN': 1,      # Complex add with negation
    'SHUFFLE': 1,    # Shuffle bytes
    'CLZ': 1,        # Count leading zeros
    'CMPSWAP': 1,    # Compare and swap
    'POPCNT': 1,     # Population count
    'SUBC': 1,       # Subtract with carry
    'SUBX': 1,       # Subtract with extended precision
    'ADDC': 1,       # Add with carry
    'ADDX': 1,       # Add with extended precision
    'CRC32B': 3,     # CRC32 byte
    'SWAPMSK': 2,    # Swap mask
    'FMAX': 4,       # Maximum of floating-point numbers
    'FLOOR': 5,      # Floor value of floating-point number
    'TRUNC': 4,      # Truncate floating-point number
    'FMOD': 6,       # Floating-point modulo
    'FMAXF': 4,      # Maximum of single-precision floating-point numbers
    'FMINF': 4,      # Minimum of single-precision floating-point numbers
    'SQRTF': 5,      # Square root of single-precision floating-point number
    'FLOORF': 5,     # Floor value of single-precision floating-point number
    'FMODF': 6,      # Floating-point modulo of single-precision number
}


def get_clock_cycles(instruction):
    """
    根据指令获取其时钟周期数。

    Args:
        instruction (str): 汇编指令字符串。

    Returns:
        int or None: 指令对应的时钟周期数，如果指令未定义则返回None。
    """
    return instruction_cycles.get(instruction.upper(), None)

def process_asm_file(filepath):
    process_id = current_process().pid
    temp_log_filename = os.path.join(os.getcwd(), f'process_asm_{process_id}.log')
    logging.basicConfig(filename=temp_log_filename, level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

    results = []
    undefined_instructions = []

    with open(filepath, 'r') as file:
        lines = file.readlines()

    filename = os.path.basename(filepath)
    current_function = None
    in_function = False
    in_section = False

    for line_num, line in enumerate(lines, 1):
        line = line.strip()

        # 匹配 .section 段，确保 .* 部分有字符
        section_match = re.match(r'^\.section\s+.+', line)
        # 修改正则表达式以捕获可能包含各种后缀的函数名
        type_match = re.match(r'^\.type\s+(\w+(\.[\w\.]*)*),\s*@function', line)
        lfe_match = re.match(r'^\.LFE', line)
        # 匹配包括标签，例如函数名:、.LFB2:、.L8:、.LVL6:、ABS20ms_enter_internal_calc.isra.0:
        label_match = re.match(r'^\w+(\.[\w\.]*)*:$', line)
        instruction_match = re.match(r'^\s*(\w+)', line)

        if section_match:
            in_section = True
            in_function = False
            current_function = None
            logging.info(f'Matched section at line {line_num}: {line}')
            continue

        if in_section:
            if type_match:
                full_function_name = type_match.group(1)
                # 去掉函数名的后缀
                current_function = re.match(r'(\w+)', full_function_name).group(1)
                in_function = True
                logging.info(f'Matched function type at line {line_num}: {line}')
                continue

            if lfe_match and in_function:
                in_function = False
                logging.info(f'Matched function end at line {line_num}: {line}')
                continue

            if label_match and in_function:
                logging.info(f'Ignoring label at line {line_num}: {line}')
                continue

            if instruction_match and in_function:
                instruction = instruction_match.group(1)
                if instruction.lower() != '.loc':
                    clock_cycles = get_clock_cycles(instruction)
                    if clock_cycles is None:
                        undefined_instructions.append(f'{instruction} at line {line_num}')
                        clock_cycles = 0
                    results.append([filename, current_function, line, clock_cycles, ''])
                continue

    if current_function:
        results[-1][4] = 0  # 在最后一个函数行重置总时钟周期

    if undefined_instructions:
        logging.warning(f'File {filename} contains undefined instructions: {set(undefined_instructions)}')

    logging.info(f'Processing file: {filename}')
    logging.shutdown()
    return results, temp_log_filename

def save_with_progress_bar(workbook, filename):
    """
    使用进度条保存Excel工作簿。

    Args:
        workbook (openpyxl.Workbook): 要保存的Excel工作簿对象。
        filename (str): 要保存的文件名。
    """
    temp_filename = filename + '.tmp'  # 临时文件名
    workbook.save(temp_filename)  # 保存到临时文件

    total_size = os.path.getsize(temp_filename)  # 计算文件总大小

    with open(temp_filename, 'rb') as temp_file, open(filename, 'wb') as final_file:
        with tqdm(total=total_size, desc='Saving to Excel', unit='B', unit_scale=True) as pbar:
            while True:
                chunk = temp_file.read(1024 * 1024)
                if not chunk:
                    break
                final_file.write(chunk)
                pbar.update(len(chunk))  # 更新进度条

    os.remove(temp_filename)  # 删除临时文件

def delete_tmp_log_files():
    # 获取当前目录下所有的文件和文件夹列表
    files = os.listdir('.')
    
    # 遍历文件列表，删除所有的 .tmp 和 .log 文件
    for file in files:
        if file.endswith('.tmp') or file.endswith('.log'):
            try:
                os.remove(file)
            except Exception as e:
                print(f"Failed to delete file: {file} - {e}")

def main():
    """
    主函数，用于处理汇编文件并将结果保存到Excel文件中。
    """
    folder_path = r'E:\eclipse-workspace\IBS_E001\E001_IBS1.1_TC377\C04_Software\06_Dev\01_Source\03_APP\Compile_Make'

    # 调用函数删除 .tmp 和 .log 文件
    delete_tmp_log_files()

    wb = Workbook()  # 创建新的Excel工作簿对象
    ws = wb.active  # 获取活动工作表对象
    ws.title = "Assembly Analysis"  # 设置工作表标题
    header = ['File Name', 'Function Name', 'Instruction', 'Clock Cycles', 'Total Cycles']
    ws.append(header)  # 添加表头行

    header_fill = PatternFill(start_color="D3D3D3", end_color="D3D3D3", fill_type="solid")
    for cell in ws[1]:
        cell.fill = header_fill  # 设置表头填充颜色

    asm_files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith('.s')]  # 获取所有汇编文件的路径列表

    with Pool(cpu_count()) as pool:
        results_and_logs = list(tqdm(pool.imap_unordered(process_asm_file, asm_files), total=len(asm_files), desc='Processing files'))

    # 合并所有子进程日志文件
    with open(log_filename, 'a') as log_file:
        for _, temp_log_filename in results_and_logs:
            if os.path.exists(temp_log_filename):
                with open(temp_log_filename, 'r') as lf:
                    log_file.write(lf.read())
                os.remove(temp_log_filename)  # 删除子进程日志文件

    # 将处理结果批量写入Excel文件
    all_results = []
    for results, _ in results_and_logs:
        all_results.extend(results)

    # 遍历并更新每个函数名的总时钟周期数
    current_function = None
    total_cycles = 0
    last_instruction_idx = None
    function_counts = {}  # 统计每个函数的指令数
    with tqdm(total=len(all_results), desc='Updating total cycles') as pbar:
        for idx, result in enumerate(all_results):
            if result[1] != current_function:
                # 计算上一个函数名的总时钟周期数并更新
                if current_function is not None and last_instruction_idx is not None:
                    logging.info(f'Updating total cycles for function {current_function}: {total_cycles}')
                    all_results[last_instruction_idx][4] = total_cycles
                    pbar.update(function_counts[current_function])  # 更新进度条
                # 开始处理新的函数名
                current_function = result[1]
                total_cycles = 0
                last_instruction_idx = None
                if current_function not in function_counts:
                    function_counts[current_function] = 0

            # 累加当前函数名的总时钟周期数
            if result[3] != '':
                total_cycles += int(result[3])
                last_instruction_idx = idx
                function_counts[current_function] += 1

        # 更新最后一个函数名的总时钟周期数
        if current_function is not None and last_instruction_idx is not None:
            logging.info(f'Updating total cycles for function {current_function}: {total_cycles}')
            all_results[last_instruction_idx][4] = total_cycles
            pbar.update(function_counts[current_function])  # 确保最后一个进度条更新到100%

    for result in tqdm(all_results, desc='Writing rows to Excel'):
        ws.append(result)

    save_with_progress_bar(wb, xlsx_filename)  # 使用带进度条的方式保存Excel文件
    logging.info(f'Processing complete. Results saved to {xlsx_filename}')  # 记录处理完成信息

if __name__ == "__main__":
    main()  # 调用主函数开始执行
